package com.cfp4cloud.cfp.knowledge.support.handler.rag.strategy;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.date.DateUtil;
import cn.hutool.core.util.BooleanUtil;
import cn.hutool.core.util.NumberUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.extra.spring.SpringUtil;
import com.cfp4cloud.cfp.common.core.constant.enums.YesNoEnum;
import com.cfp4cloud.cfp.common.milvus.MilvusEmbeddingStore;
import com.cfp4cloud.cfp.knowledge.config.properties.AiKnowledgeProperties;
import com.cfp4cloud.cfp.knowledge.config.properties.FullTextSearchProperties;
import com.cfp4cloud.cfp.knowledge.dto.AiMessageResultDTO;
import com.cfp4cloud.cfp.knowledge.dto.AiSliceReCallRequestDTO;
import com.cfp4cloud.cfp.knowledge.dto.ChatMessageDTO;
import com.cfp4cloud.cfp.knowledge.entity.AiChatRecordEntity;
import com.cfp4cloud.cfp.knowledge.entity.AiDatasetEntity;
import com.cfp4cloud.cfp.knowledge.entity.AiDocumentEntity;
import com.cfp4cloud.cfp.knowledge.entity.AiSliceEntity;
import com.cfp4cloud.cfp.knowledge.service.AiDocumentService;
import com.cfp4cloud.cfp.knowledge.service.AiStreamAssistantService;
import com.cfp4cloud.cfp.knowledge.support.constant.DocumentTypeEnums;
import com.cfp4cloud.cfp.knowledge.support.constant.SourceTypeEnums;
import com.cfp4cloud.cfp.knowledge.support.constant.SummaryStatusEnums;
import com.cfp4cloud.cfp.knowledge.support.provider.ModelProvider;
import com.cfp4cloud.cfp.knowledge.support.util.PromptBuilder;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.scoring.ScoringModel;
import dev.langchain4j.rag.content.Content;
import dev.langchain4j.rag.content.aggregator.ContentAggregator;
import dev.langchain4j.rag.content.aggregator.DefaultContentAggregator;
import dev.langchain4j.rag.content.aggregator.ReRankingContentAggregator;
import dev.langchain4j.rag.query.Query;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingSearchResult;
import dev.langchain4j.store.embedding.EmbeddingStore;
import lombok.SneakyThrows;
import lombok.experimental.UtilityClass;
import lombok.extern.slf4j.Slf4j;
import reactor.core.publisher.Flux;

import java.util.*;

import static com.cfp4cloud.cfp.knowledge.support.constant.AiPromptField.systemTime;
import static dev.langchain4j.store.embedding.filter.MetadataFilterBuilder.metadataKey;
import static java.util.Collections.singletonList;
import static java.util.Collections.singletonMap;

/**
 * 统一RAG辅助工具类
 * <p>
 * 合并了原RAG处理和切片向量化的功能，提供完整的RAG处理辅助方法： 1. 向量搜索：构建搜索请求、执行检索 2. 全文检索：BM25搜索支持 3.
 * 重排序：基于重排序模型的结果优化 4. 答案生成：AI模型结果总结 5. 数据转换：DTO转换、链接构建 6. 工具方法：空值判断、ID提取等
 * </p>
 *
 * @author pig
 * @date 2024-03-14
 */
@Slf4j
@UtilityClass
public class RagHelper {

	/**
	 * 默认内容聚合器
	 */
	public static final ContentAggregator DEFAULT_AGGREGATOR = new DefaultContentAggregator();

	/**
	 * 构建向量搜索请求
	 * @param queryEmbedding 查询向量
	 * @param dataset 数据集配置
	 * @param documentType 文档类型
	 * @param minScore 最小分数（可选）
	 * @return 搜索请求对象
	 */
	public static EmbeddingSearchRequest buildSearchRequest(Embedding queryEmbedding, AiDatasetEntity dataset,
			String documentType, Double minScore) {
		double finalMinScore;
		if (Objects.nonNull(minScore)) {
			finalMinScore = minScore;
		}
		else {
			// 使用数据集配置的默认分数阈值
			finalMinScore = NumberUtil.div(Double.parseDouble(dataset.getScore().toString()), 100.0, 2);
		}

		return new EmbeddingSearchRequest(queryEmbedding, dataset.getTopK(), finalMinScore,
				metadataKey(AiDocumentEntity.Fields.datasetId).isEqualTo(dataset.getId().toString())
					.and(metadataKey(DocumentTypeEnums.Fields.type).isEqualTo(documentType)));
	}

	/**
	 * 执行全文检索（仅支持 Milvus）
	 * @param embeddingStore 向量存储
	 * @param queryContent 查询内容
	 * @return 全文检索结果
	 */
	public static List<EmbeddingMatch<TextSegment>> performFullTextSearch(EmbeddingStore<TextSegment> embeddingStore,
			String queryContent) {
		if (!(embeddingStore instanceof MilvusEmbeddingStore milvusEmbeddingStore)) {
			return Collections.emptyList();
		}

		try {
			AiKnowledgeProperties knowledgeProperties = SpringUtil.getBean(AiKnowledgeProperties.class);
			FullTextSearchProperties fullTextSearch = knowledgeProperties.getFullTextSearch();

			if (!fullTextSearch.isEnabled()) {
				return Collections.emptyList();
			}

			Map<String, Object> searchParams = Map.of(fullTextSearch.getAlgorithm(), fullTextSearch.getDropRatio());
			EmbeddingSearchResult<TextSegment> searchResult = milvusEmbeddingStore.searchByBm25(queryContent,
					searchParams, fullTextSearch.getTopK());

			if (Objects.nonNull(searchResult) && !searchResult.matches().isEmpty()) {
				return searchResult.matches();
			}
		}
		catch (Exception e) {
			log.warn("全文检索失败: {}", e.getMessage(), e);
		}

		return Collections.emptyList();
	}

	/**
	 * 对搜索结果进行重排序
	 * @param dataset 数据集配置
	 * @param queryContent 查询内容
	 * @param embeddingMatches 匹配结果
	 * @return 重排序后的内容列表
	 */
	public static List<Content> rerankSearchResults(AiDatasetEntity dataset, String queryContent,
			List<EmbeddingMatch<TextSegment>> embeddingMatches) {
		Query query = Query.from(queryContent);
		Map<Query, Collection<List<Content>>> queryToContents = singletonMap(query,
				singletonList(embeddingMatches.stream().map(match -> Content.from(match.embedded().text())).toList()));

		// 如果没有配置重排序模型，使用默认聚合器
		if (StrUtil.isBlank(dataset.getRerankerModel())) {
			return DEFAULT_AGGREGATOR.aggregate(queryToContents);
		}

		try {
			ModelProvider modelProvider = SpringUtil.getBean(ModelProvider.class);
			ScoringModel rerankerModel = modelProvider.getRerankerModel(dataset.getRerankerModel());

			if (Objects.nonNull(rerankerModel)) {
				ContentAggregator contentAggregator = ReRankingContentAggregator.builder()
					.scoringModel(rerankerModel)
					.build();
				return contentAggregator.aggregate(queryToContents);
			}
		}
		catch (Exception e) {
			log.warn("重排序失败: {}", e.getMessage(), e);
		}

		return DEFAULT_AGGREGATOR.aggregate(queryToContents);
	}

	/**
	 * 处理空搜索结果
	 * @param dataset 数据集配置
	 * @param chatMessageDTO 聊天消息
	 * @return 空结果响应
	 */
	public static Flux<AiMessageResultDTO> handleEmptyResult(AiDatasetEntity dataset, ChatMessageDTO chatMessageDTO) {
		// 根据配置决定返回固定回复或使用大模型生成回答
		if (YesNoEnum.NO.getCode().equals(dataset.getEmptyLlmFlag())) {
			// 直接返回预设的空结果描述
			return Flux.just(new AiMessageResultDTO(dataset.getEmptyDesc()));
		}
		else {
			// 使用大模型生成回答
			try {
				ModelProvider modelProvider = SpringUtil.getBean(ModelProvider.class);
				AiStreamAssistantService streamAssistantService = modelProvider
					.getAiStreamAssistant(chatMessageDTO.getModelName())
					.getValue();

				return streamAssistantService.chat(chatMessageDTO.getConversationId(), chatMessageDTO.getContent())
					.map(text -> {
						AiMessageResultDTO dto = new AiMessageResultDTO(text);
						dto.setExtLinks(List.of(new AiMessageResultDTO.ExtLink("未匹配知识库内容，由大模型直接推理返回", null, null)));
						return dto;
					});
			}
			catch (Exception e) {
				log.error("大模型回答生成失败: {}", e.getMessage(), e);
				return Flux.just(new AiMessageResultDTO(dataset.getEmptyDesc()));
			}
		}
	}

	/**
	 * 构建参考资料链接
	 * @param embeddingMatches 匹配结果
	 * @param documentService 文档服务
	 * @return 参考链接列表
	 */
	@SneakyThrows
	public static List<AiMessageResultDTO.ExtLink> buildReferenceLinks(
			List<EmbeddingMatch<TextSegment>> embeddingMatches, AiDocumentService documentService) {
		List<AiMessageResultDTO.ExtLink> extLinks = new ArrayList<>();

		for (EmbeddingMatch<TextSegment> match : embeddingMatches) {
			String documentId = match.embedded().metadata().getString(AiSliceEntity.Fields.documentId);
			if (Objects.isNull(documentId)) {
				continue;
			}

			Float distance = match.score().floatValue();
			AiDocumentEntity document = documentService.getById(Long.parseLong(documentId));
			if (Objects.nonNull(document)) {
				extLinks.add(new AiMessageResultDTO.ExtLink(document.getName(),
						String.format("/admin/sys-file/oss/file?fileName=%s", document.getFileUrl()), distance));
			}
		}

		return extLinks.stream().distinct().toList();
	}

	/**
	 * 提取标准问答的答案
	 * @param embeddingMatches 匹配结果
	 * @return 答案列表
	 */
	public static List<String> extractStandardAnswers(List<EmbeddingMatch<TextSegment>> embeddingMatches) {
		return embeddingMatches.stream()
			.map(match -> match.embedded().metadata().getString(AiChatRecordEntity.Fields.answerText))
			.filter(Objects::nonNull)
			.toList();
	}

	/**
	 * 合并多个搜索结果（支持可变参数）
	 * @param searchResults 多个搜索结果列表
	 * @return 合并后的结果
	 */
	@SafeVarargs
	public static List<EmbeddingMatch<TextSegment>> mergeSearchResults(
			List<EmbeddingMatch<TextSegment>>... searchResults) {
		List<EmbeddingMatch<TextSegment>> mergedResults = new ArrayList<>();

		for (List<EmbeddingMatch<TextSegment>> results : searchResults) {
			if (CollUtil.isNotEmpty(results)) {
				mergedResults.addAll(results);
			}
		}

		return mergedResults;
	}

	/**
	 * 验证搜索结果是否为空
	 * @param embeddingMatches 搜索结果
	 * @return true表示为空，false表示有结果
	 */
	public static boolean isEmpty(List<EmbeddingMatch<TextSegment>> embeddingMatches) {
		return CollUtil.isEmpty(embeddingMatches);
	}

	/**
	 * 提取嵌入ID列表（用于统计命中次数）
	 * @param embeddingMatches 匹配结果
	 * @return 嵌入ID列表
	 */
	public static List<String> extractEmbeddingIds(List<EmbeddingMatch<TextSegment>> embeddingMatches) {
		return embeddingMatches.stream().map(EmbeddingMatch::embeddingId).filter(Objects::nonNull).toList();
	}

	/**
	 * 结果总结输出 使用AI模型对检索到的多个结果进行总结，生成最终回答
	 * @param dataset 数据集配置信息
	 * @param chatMessageDTO 聊天消息DTO，包含用户问题和对话上下文
	 * @param resultList 从向量库检索到的相关文本片段列表
	 * @return 返回AI生成的总结回答流
	 */
	public static Flux<AiMessageResultDTO> summaryResult(AiDatasetEntity dataset, ChatMessageDTO chatMessageDTO,
			List<String> resultList) {
		// 对结果进行总结
		ModelProvider modelProvider = SpringUtil.getBean(ModelProvider.class);
		AiStreamAssistantService streamAssistantService = modelProvider
			.getAiStreamAssistant(chatMessageDTO.getModelName())
			.getValue();
		Flux<String> summaryResult = streamAssistantService.chat(chatMessageDTO.getConversationId(),
				PromptBuilder.render("knowledge-system.st"),
				PromptBuilder.render("knowledge-user.st",
						Map.of("contents", CollUtil.join(resultList, StrUtil.CRLF), "userMessage",
								chatMessageDTO.getContent(), "emptyDesc", dataset.getEmptyDesc(), systemTime,
								DateUtil.now())));
		return summaryResult.map(AiMessageResultDTO::new);
	}

	// ========== 以下为切片向量化相关方法 ==========

	/**
	 * 判断是否应跳过摘要处理
	 * @param aiDataset 知识库配置
	 * @param documentEntity 文档实体
	 * @return true表示应跳过，false表示继续处理
	 */
	public static boolean shouldSkipSummary(AiDatasetEntity aiDataset, AiDocumentEntity documentEntity) {
		return YesNoEnum.YES.getCode().equals(aiDataset.getPreSummary())
				&& !SourceTypeEnums.QA.getType().equals(documentEntity.getSourceType())
				&& !documentEntity.getSummaryStatus().equals(SummaryStatusEnums.SUMMARYED.getStatus());
	}

	/**
	 * 构建用于切片召回的向量搜索请求（重载方法，支持不同参数）
	 * @param queryEmbedding 查询向量
	 * @param requestDTO 搜索请求参数
	 * @param dataset 数据集配置
	 * @return 向量搜索请求对象
	 */
	public static EmbeddingSearchRequest buildSearchRequest(Embedding queryEmbedding,
			AiSliceReCallRequestDTO requestDTO, AiDatasetEntity dataset) {
		double minScore = NumberUtil.div(Double.parseDouble(dataset.getScore().toString()), 100.0, 2);

		int maxResults = Objects.nonNull(requestDTO.getTopK()) ? requestDTO.getTopK() : dataset.getTopK();

		return EmbeddingSearchRequest.builder()
			.queryEmbedding(queryEmbedding)
			.maxResults(maxResults)
			.filter(metadataKey(AiDocumentEntity.Fields.datasetId).isEqualTo(dataset.getId().toString())
				.and(metadataKey(DocumentTypeEnums.Fields.type).isEqualTo(DocumentTypeEnums.ANSWER.getType())))
			.minScore(minScore)
			.build();
	}

	/**
	 * 执行全文检索（支持切片召回请求）
	 * @param embeddingStore 向量存储
	 * @param requestDTO 请求参数
	 * @return 全文检索结果，如果不支持则返回空列表
	 */
	public static List<EmbeddingMatch<TextSegment>> performFullTextSearch(EmbeddingStore<TextSegment> embeddingStore,
			AiSliceReCallRequestDTO requestDTO) {
		if (!(embeddingStore instanceof MilvusEmbeddingStore milvusEmbeddingStore)
				|| !BooleanUtil.toBoolean(requestDTO.getFullTextSearch())) {
			return Collections.emptyList();
		}

		try {
			AiKnowledgeProperties knowledgeProperties = SpringUtil.getBean(AiKnowledgeProperties.class);
			FullTextSearchProperties fullTextSearch = knowledgeProperties.getFullTextSearch();

			// 配置BM25算法参数
			Map<String, Object> searchParams = Map.of(fullTextSearch.getAlgorithm(), fullTextSearch.getDropRatio());

			// 执行BM25全文检索
			EmbeddingSearchResult<TextSegment> searchResult = milvusEmbeddingStore.searchByBm25(requestDTO.getContent(),
					searchParams, fullTextSearch.getTopK());

			if (Objects.nonNull(searchResult) && !searchResult.matches().isEmpty()) {
				return searchResult.matches();
			}
		}
		catch (Exception e) {
			log.warn("全文检索失败: {}", e.getMessage(), e);
		}

		return Collections.emptyList();
	}

	/**
	 * 执行重排序（简化版，支持查询内容）
	 * @param embeddingMatches 原始检索结果
	 * @param queryContent 查询内容
	 * @return 重排序后的结果
	 */
	public static List<EmbeddingMatch<TextSegment>> performReranking(List<EmbeddingMatch<TextSegment>> embeddingMatches,
			String queryContent) {
		if (embeddingMatches.isEmpty()) {
			return embeddingMatches;
		}

		try {
			// 使用默认聚合器进行重排序
			Query query = Query.from(queryContent);
			Map<Query, Collection<List<Content>>> queryToContents = singletonMap(query, singletonList(
					embeddingMatches.stream().map(match -> Content.from(match.embedded().text())).toList()));

			return DEFAULT_AGGREGATOR.aggregate(queryToContents)
				.stream()
				.map(content -> embeddingMatches.stream()
					.filter(match -> content.textSegment().text().equals(match.embedded().text()))
					.findFirst()
					.orElse(null))
				.filter(Objects::nonNull)
				.toList();
		}
		catch (Exception e) {
			log.warn("重排序失败: {}", e.getMessage(), e);
			return embeddingMatches;
		}
	}

}